import argparse

import torch

def parse_args():
    parser = argparse.ArgumentParser()

    # message
    parser.add_argument('--msg_batch', action='store_true')
    parser.add_argument('--msg_epoch', action='store_false')
    parser.add_argument('--msg_batch_freq', type=int, default=1000)

    # log
    parser.add_argument('--log_name', default='log')
    parser.add_argument('--log_epoch_freq', type=int, default=1)
    parser.add_argument('--log_batch_freq', type=int, default=0)

    # data
    parser.add_argument('--batch', type=int, default=4)
    parser.add_argument('--workers', type=int, default=0)
    parser.add_argument('--epoch', type=int, default=10)

    # model
    parser.add_argument('--model', default='resnet')
    parser.add_argument('--size', type=int, default=480)
    parser.add_argument('--classes', type=int, default=100)
    parser.add_argument('--load')
    parser.add_argument('--torch_pretrain', action='store_true')
    parser.add_argument('--save_dir')
    parser.add_argument('--save_freq', type=int, default=10)

    # eval
    parser.add_argument('--eval_batch_freq', type=int, default=0)
    parser.add_argument('--eval_epoch_freq', type=int, default=1)

    # scheduler
    parser.add_argument('--scheduler', default='none')
    parser.add_argument('--milestones', nargs='+', type=int)
    parser.add_argument('--gamma', type=float, default=1e-1)

    # optimizer
    parser.add_argument('--optim', default='Adam')
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--betas', nargs='+', type=float, default=[0.95, 0.999])
    parser.add_argument('--weight_decay', type=float, default=0)
    parser.add_argument('--momentum', type=float, default=0.9)

    # resume
    parser.add_argument('--resume', action='store_true')

    # GPU
    parser.add_argument('--device', default='cpu')

    return parser.parse_args()

opt = parse_args()

if 'cuda' in opt.device:
    if not torch.cuda.is_available():
        opt.device = 'cpu'